#!/usr/bin/env python3
# -*- coding: utf-8 -*-
################################
# File Name   : smart_link
# Author      : liyanqing
# Created On  : 2020-09-16 07:11:28
# Description : link A to B, keep dir and link file.
################################
import os
import re
import sys
import shutil
import argparse

os.environ['PYTHONUNBUFFERED'] = '1'


def read_args():
    """
    Read in arguments.
    """
    parser = argparse.ArgumentParser()

    parser.add_argument('source_path',
                        help='Specify source path.')
    parser.add_argument('target_path',
                        help='Specify target path.')
    parser.add_argument('-o', '--overwrite',
                        action='store_true',
                        default=False,
                        help='Overwrite by force if target_path exists.')

    args = parser.parse_args()

    # Check source_path exists or not.
    if not os.path.exists(args.source_path):
        print('*Error*: "' + str(args.source_path) + '": No such file or directory!')
        sys.exit(1)

    # Remove target_path if it exists on overwrite mode.
    if os.path.exists(args.target_path):
        if not args.overwrite:
            print('*Error*: target path "' + str(args.target_path) + '" exists, please remove it first, or enable overwrite mode.')
            sys.exit(1)
        else:
            remove_path(args.target_path)

    args.target_path = os.path.abspath(args.target_path)

    return args.source_path, args.target_path


def create_directory(dir_path):
    """
    Create directory with os.makedirs, just like "mkdir".
    """
    if not os.path.exists(dir_path):
        print('mkdir ' + str(dir_path))

        try:
            os.makedirs(dir_path)
        except Exception as error:
            print('*Error*: Failed on creating directory "' + str(dir_path) + '", ' + str(error))
            sys.exit(1)


def remove_path(item_path):
    """
    Remove file/directory with shutil.rmtree/os.remove, just like "rm -rf".
    """
    if os.path.exists(item_path):
        print('*Warning*: file/directory "' + str(item_path) + '" exists, removing ...')

    print('rm -rf ' + str(item_path))

    try:
        shutil.rmtree(item_path)
    except Exception as error:
        if os.path.isfile(item_path) or os.path.islink(item_path):
            try:
                os.remove(item_path)
            except Exception as error:
                print('*Error*: Failed on removing file/directory "' + str(item_path) + '", ' + str(error))
                sys.exit(1)
        else:
            print('*Error*: Failed on removing file/directory "' + str(item_path) + '", ' + str(error))
            sys.exit(1)


def dir_link(source_dir, target_dir):
    """
    Copy source_dir to target_dir with mkdir/link.
    """
    if not os.path.exists(source_dir):
        print('*Warning*: source directory "' + str(source_dir) + '" is missing.')
    else:
        if not os.path.exists(target_dir):
            create_directory(target_dir)

        cwd = os.getcwd()

        # Come into target_dir.
        print('cd ' + str(target_dir))
        os.chdir(target_dir)

        # Get real source_dir path for relative format.
        if not re.match(r'^/.*$', source_dir):
            source_dir = os.path.join('..', source_dir)

        # Process sub file/dir.
        for sub_item in os.listdir(source_dir):
            source_dir_tmp = os.path.join(source_dir, sub_item)

            # Remove exists sub file/dir.
            if os.path.exists(sub_item):
                remove_path(sub_item)

            # Iteration link.
            if os.path.isdir(source_dir_tmp):
                dir_link(source_dir_tmp, sub_item)
            else:
                file_link(source_dir_tmp, sub_item)

        # Come back.
        print('cd ' + str(cwd))
        os.chdir(cwd)


def file_link(source_file, target_file):
    """
    Link source_file to target_file with os.symlink, just like "ln -s".
    """
    if not os.path.exists(source_file):
        print('*Warning*: source file "' + str(source_file) + '" is missing.')
    else:
        if os.path.exists(target_file):
            print('*Warning*: target file "' + str(target_file) + '" exists.')
        else:
            print('ln -s ' + str(source_file) + ' ' + str(target_file))

            try:
                os.symlink(source_file, target_file)
            except Exception as error:
                print('*Error*: Failed on linking "' + str(source_file) + '" to "' + str(target_file) + '", ' + str(error))
                sys.exit(1)


def smart_link(source_path, target_path):
    """
    Link directory with dir_link().
    Link file with file_link().
    """
    if os.path.isdir(source_path):
        dir_link(source_path, target_path)
    else:
        file_link(source_path, target_path)


################
# Main Process #
################
def main():
    (source_path, target_path) = read_args()
    smart_link(source_path, target_path)


if __name__ == '__main__':
    main()
